use cudarc::driver::{LaunchConfig, PushKernelArg};

fn main() -> Result<(), cudarc::driver::result::DriverError> {
    // Get a context for GPU 0
    let ctx = cudarc::driver::CudaContext::new(0)?;
    let stream = ctx.default_stream();

    // Copy data to device and allocate output buffer
    let inp = stream.memcpy_stod(&[1.0f32; 100])?;
    let mut out = stream.alloc_zeros::<f32>(100)?;

    // Compile and load a kernel
    let ptx = cudarc::nvrtc::compile_ptx(
        r#"
extern "C" __global__ 
void my_kernel(float* out, float* in, int num_elem)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x; 
    if (idx < num_elem) 
    {
        out[idx] = in[idx] * 2.0 + idx;
    }
}
    "#,
    )
    .unwrap();
    let module = ctx.load_module(ptx)?;
    let kernel = module.load_function("my_kernel")?;

    // Launch the kernel
    let mut builder = stream.launch_builder(&kernel);
    builder.arg(&mut out);
    builder.arg(&inp);
    builder.arg(&100usize);
    unsafe { builder.launch(LaunchConfig::for_num_elems(100)) }?;

    // Copy results back to host
    let results: Vec<f32> = stream.memcpy_dtov(&out)?;

    println!("result: {results:?}");

    Ok(())
}
